package leetcode_2101_2200;
import java.util.Arrays;
import java.util.TreeSet;

public class LeeCode_2195_wrong {
    public static void main(String[] args) {
//        System.out.println(minimalKSum(new int[]{1, 4, 25, 10, 15}, 2));
//        System.out.println(minimalKSum(new int[]{5,6}, 3));
//        System.out.println(minimalKSum(new int[]
//                {96,44,99,25,61,84,88,18,19,33,60,86,52,19,32,47,35,50,94,17,29,98,22,21,72,100,40,84}, 35));
        System.out.println(getSum(1, 3));
    }
    private static long minimalKSum(int[] nums, int k) {
        TreeSet<Integer> set = new TreeSet<>();
        for (int num : nums) {
            set.add(num);
        }
        nums = new int[set.size()];
        int tmp = 0;
        while (tmp < nums.length){
            int num = set.pollFirst();
            nums[tmp] = num;
            tmp++;
        }
        int n = nums.length;
        int index = bigger(nums, k);
        long ans = 0;
        if (index == -1){
            long sum = Arrays.stream(nums).sum();
            long all = getSum(1, nums[n - 1]);
            ans += all - sum;
            int remain = k - (nums[n - 1] - (n - 1) - 1);
            int l = nums[n - 1] + 1, r = l + remain - 1;
            ans += getSum(l, r);
            return ans;
        }
        if (index == 0){
            ans += getSum(1, k);
            return ans;
        }
        long sum = 0;
        for (int i = 0; i < index; i++) {
            sum += nums[i];
        }
        ans += getSum(1, nums[index - 1]) - sum;
        int remain = k - (nums[index - 1] - (index - 1) - 1);
        int l = nums[index - 1] + 1, r = l + remain - 1;
        ans += getSum(l, r);
        return ans;
    }
    private static long getSum(int l, int r){
        long ans = (long) (l + r) * ((r - l + 1) / 2);
        if (((r - l + 1 ) & 1) == 1){
            ans += (l + r) / 2;
        }
        return ans;
    }
    private static int bigger(int[] nums, int k){
        long l = 0, r = nums.length - 1;
        while (l < r){
            long m = (l + r) / 2;
            long forget = nums[(int) m] - m - 1;
            if (forget < k){
                l = m + 1;
            }else {
                r = m;
            }
        }
        if (nums[(int) l] - l - 1 < k)
            return -1;
        return (int) l;
    }
}